# test.py文件
import torch
from torchvision import transforms
from PIL import Image

数据预处理 = transforms.Compose([
    transforms.Resize(256),  # 将图像大小调整为 256x256
    transforms.CenterCrop(224),  # 从中心裁剪出 224x224 的图像
    transforms.ToTensor(),  # 将图像转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 归一化
])
# 加载模型
模型 = torch.load('模型文件.pth', weights_only=False)
模型.eval()  # 将模型设置为评估模式
# 指定设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
模型.to(device)
# 读取图像
image_path = '1.jpg'  # 替换为你的图像路径
image_paths = ['键盘.jpg', '杯子.jpg']
for image_path in image_paths:
    image = Image.open(image_path).convert('RGB')  # 确保图像是RGB格式
    # 预处理图像
    image = 数据预处理(image).unsqueeze(0)  # 添加批量维度
    image = image.to(device)
    # 进行预测
    with torch.no_grad():
        output = 模型(image)
        _, predicted = torch.max(output, 1)
    # 获取类别名称
    # CIFAR-100 的类别名称
    class_names = [
        "苹果", "水族馆鱼", "婴儿", "熊", "海狸", "床", "蜜蜂", "甲虫",
        "自行车", "瓶子", "碗", "男孩", "桥梁", "公共汽车", "蝴蝶", "骆驼",
        "罐子", "城堡", "毛毛虫", "家畜", "椅子", "黑猩猩", "时钟",
        "云", "蟑螂", "沙发", "螃蟹", "鳄鱼", "杯子", "恐龙",
        "海豚", "大象", "比目鱼", "森林", "狐狸", "女孩", "仓鼠",
        "房子", "袋鼠", "键盘", "台灯", "割草机", "豹", "狮子",
        "蜥蜴", "龙虾", "男人", "枫树", "摩托车", "山", "老鼠",
        "蘑菇", "橡树", "橘子", "兰花", "水獭", "棕榈树", "梨",
        "卡车", "火车", "柳树", "狼", "女人", "鲸鱼", "斑马"
    ]
    predicted_class = class_names[predicted.item()]
    print(f'图片识别结果是: {predicted_class}')
